#!/usr/bin/env python3
'''
This class is used to extract data from the file describing geometry of the problem.
'''

import re
import numpy


class Hamiltonian:
    def __init__(self, hamiltonianText):
        self.hamiltonian = hamiltonianText

        self.hamiltonian.strip().split('\n')
        self.t = {} #Hopping part
        self.u = {} #Potential energy part

        for line in self.hamiltonian:
            temp = line.split()

            if float(temp[-1]) == 0:#hopping term
                hop = Hop()
                hop.x = float(temp[2])
                hop.y = float(temp[3])
                hop.z = float(temp[4])
                hop.t_up = float(temp[5])
                hop.t_down = float(temp[6])

            else:#potential energy term
                self.u[int(temp[0])] = float(temp[-1])


class Hop:
    pass


class Orbits:
    def __init__(self, orbitText):
        self.orbits = {}
        for lineIndex, line in enumerate(orbitText.strip().split("\n")):
            orb = Orbit()
            orb.index = lineIndex
            temp = line.split()
            orb.x = float(temp[1].replace("d0", ""))
            orb.y = float(temp[2].replace("d0", ""))
            orb.z = float(temp[3].replace("d0", ""))
            self.orbits[lineIndex] = orb
        self.nOrbits = len(self.orbits)

    def __iter__(self):
        return self.orbits


class Orbit:
    pass

class Parser:
    def __init__(self, geometryTextFile):
        self.dimension = int(re.search('(?<=NDIM)\s\d', geometryTextFile).group(0))#dimension of the problem
        self.prim = re.search(r'(?<=PRIM)\s*((\d+.\d+)+\s*)', geometryTextFile).group(0)#array storing primitive vectors
        if (self.dimension == 1):
            self.prim = float(self.prim.strip())
        self.super = re.search(r'(?<=SUPER)\s*((\d+)+\s*)', geometryTextFile).group(0)#array storing primitive vectors
        if (self.dimension == 1):
            self.super = float(self.super.strip())

        self.orbits = Orbits(re.search(r'(?<=ORB)(\s*\w*\.?\w*)+',
                                      geometryTextFile).group(0))#string storing orbits

        #self.hamiltonian = Hamiltonian(re.search(r'(?<HAMILT)\.*(\d.*)\#', geometryTextFile))#description of the hoppings




